"""
模型下载组件
"""

import tkinter as tk
from tkinter import messagebox, ttk
import threading
from pathlib import Path
from typing import TYPE_CHECKING

from ...utils import FileUtils

if TYPE_CHECKING:
    from ..main_window import MainWindow


class ModelDownloadMixin:
    """模型下载 Mixin 类"""
    
    def download_model(self: 'MainWindow'):
        """下载模型对话框"""
        # 创建下载窗口
        download_window = tk.Toplevel(self.root)
        download_window.title("下载模型")
        download_window.geometry("650x500")
        download_window.transient(self.root)
        download_window.grab_set()
        
        # 说明文本
        info_frame = ttk.Frame(download_window, padding=10)
        info_frame.pack(fill=tk.X)
        
        ttk.Label(
            info_frame,
            text="选择要下载的模型类型：",
            font=("Arial", 12, "bold")
        ).pack(anchor=tk.W, pady=(0, 10))
        
        # 模型类型选择
        model_type_var = tk.StringVar(value="sam")
        
        type_frame = ttk.Frame(info_frame)
        type_frame.pack(fill=tk.X, pady=5)
        
        ttk.Radiobutton(
            type_frame,
            text="SAM 分割模型（图片分割）",
            variable=model_type_var,
            value="sam",
            command=lambda: self._update_model_selection(model_type_var.get(), model_var, models_frame)
        ).pack(anchor=tk.W)
        
        ttk.Radiobutton(
            type_frame,
            text="MiDaS 深度估计模型（景深模式）",
            variable=model_type_var,
            value="midas",
            command=lambda: self._update_model_selection(model_type_var.get(), model_var, models_frame)
        ).pack(anchor=tk.W)
        
        ttk.Radiobutton(
            type_frame,
            text="Rembg 背景移除模型（AI去背景）",
            variable=model_type_var,
            value="rembg",
            command=lambda: self._update_model_selection(model_type_var.get(), model_var, models_frame)
        ).pack(anchor=tk.W)
        
        # 分隔线
        ttk.Separator(info_frame, orient='horizontal').pack(fill=tk.X, pady=10)
        
        ttk.Label(
            info_frame,
            text="选择具体模型：",
            font=("Arial", 11, "bold")
        ).pack(anchor=tk.W)
        
        # 模型选择框架
        models_frame = ttk.Frame(info_frame)
        models_frame.pack(fill=tk.X, pady=5)
        
        # SAM模型信息
        sam_models_info = self.config.get('model.download_urls', {})
        
        model_var = tk.StringVar(value="vit_b")
        
        # 初始显示SAM模型
        self._populate_model_options(models_frame, model_var, sam_models_info, "sam")
        
        # 进度显示
        progress_frame = ttk.Frame(download_window, padding=10)
        progress_frame.pack(fill=tk.BOTH, expand=True)
        
        progress_text = tk.Text(progress_frame, height=10, width=70)
        progress_text.pack(fill=tk.BOTH, expand=True)
        
        progress_bar = ttk.Progressbar(progress_frame, mode='determinate', length=500, bootstyle="success-striped")
        progress_bar.pack(fill=tk.X, pady=5)
        
        # 按钮
        button_frame = ttk.Frame(download_window, padding=10)
        button_frame.pack(fill=tk.X)
        
        def start_download():
            model_type = model_type_var.get()
            choice = model_var.get()
            
            if model_type == "rembg":
                # Rembg模型说明（自动下载）
                messagebox.showinfo(
                    "Rembg模型下载",
                    f"Rembg模型无需手动下载！\n\n"
                    f"首次使用时会自动下载所选模型。\n\n"
                    f"使用方法：\n"
                    f"1. 确保已安装rembg库：pip install rembg\n"
                    f"2. 开启\"Rembg去背景\"开关\n"
                    f"3. 首次使用会自动下载模型\n\n"
                    f"当前选择的模型：{choice}",
                    parent=download_window
                )
                return
            
            if model_type == "midas":
                # MiDaS模型手动下载
                midas_download_urls = {
                    'DPT_Large': {
                        'url': 'https://github.com/intel-isl/DPT/releases/download/1_0/dpt_large-midas-2f21e586.pt',
                        'filename': 'dpt_large-midas-2f21e586.pt',
                        'size': '约1.4GB'
                    },
                    'DPT_Hybrid': {
                        'url': 'https://github.com/intel-isl/DPT/releases/download/1_0/dpt_hybrid-midas-501f0c75.pt',
                        'filename': 'dpt_hybrid-midas-501f0c75.pt',
                        'size': '约700MB'
                    },
                    'MiDaS_small': {
                        'url': 'https://github.com/intel-isl/MiDaS/releases/download/v2_1/midas_v21_small-70d6b9c8.pt',
                        'filename': 'midas_v21_small-70d6b9c8.pt',
                        'size': '约100MB'
                    }
                }
                
                model_info = midas_download_urls[choice]
                
                # MiDaS模型下载到torch缓存目录
                import torch
                torch_home = torch.hub.get_dir()
                checkpoints_dir = Path(torch_home) / 'checkpoints'
                checkpoints_dir.mkdir(parents=True, exist_ok=True)
                
                output_path = checkpoints_dir / model_info["filename"]
            else:
                # SAM模型下载
                model_info = sam_models_info[choice]
                
                # 创建 models 目录
                models_dir = Path(self.config.get('model.models_dir', 'models'))
                models_dir.mkdir(exist_ok=True)
                
                output_path = models_dir / model_info["filename"]
            
            # 检查文件是否已存在
            if output_path.exists():
                if not messagebox.askyesno(
                    "文件已存在",
                    f"文件已存在:\n{output_path}\n\n是否重新下载？",
                    parent=download_window
                ):
                    return
                output_path.unlink()
            
            # 禁用按钮
            download_btn.config(state='disabled')
            
            # 在后台线程下载
            def download_thread():
                try:
                    progress_text.insert(tk.END, f"开始下载: {model_info['filename']}\n")
                    progress_text.insert(tk.END, f"URL: {model_info['url']}\n")
                    progress_text.insert(tk.END, f"大小: {model_info['size']}\n\n")
                    progress_text.see(tk.END)
                    
                    def show_progress(block_num, block_size, total_size):
                        downloaded = block_num * block_size
                        percent = min(downloaded * 100.0 / total_size, 100)
                        
                        download_window.after(0, lambda: progress_bar.config(value=percent))
                        
                        if block_num % 100 == 0:
                            msg = f"下载进度: {percent:.1f}% ({downloaded / 1024 / 1024:.1f}MB / {total_size / 1024 / 1024:.1f}MB)\n"
                            download_window.after(0, lambda: progress_text.insert(tk.END, msg))
                            download_window.after(0, lambda: progress_text.see(tk.END))
                    
                    success = FileUtils.download_file(
                        model_info['url'],
                        str(output_path),
                        show_progress
                    )
                    
                    if success:
                        download_window.after(0, lambda: progress_text.insert(tk.END, "\n✓ 下载完成！\n"))
                        download_window.after(0, lambda: progress_text.see(tk.END))
                        
                        # 根据模型类型显示不同的成功消息
                        if model_type_var.get() == "midas":
                            success_msg = f"MiDaS模型下载成功！\n\n保存位置:\n{output_path}\n\n现在可以点击\"生成景深图\"使用该模型。"
                        else:
                            success_msg = f"SAM模型下载成功！\n\n保存位置:\n{output_path}\n\n现在可以点击\"加载模型\"使用该模型。"
                        
                        download_window.after(0, lambda: messagebox.showinfo("成功", success_msg, parent=download_window))
                        download_window.after(0, download_window.destroy)
                    else:
                        raise Exception("下载失败")
                    
                except Exception as e:
                    error_msg = f"\n✗ 下载失败: {str(e)}\n"
                    download_window.after(0, lambda: progress_text.insert(tk.END, error_msg))
                    download_window.after(0, lambda: progress_text.see(tk.END))
                    download_window.after(0, lambda: messagebox.showerror("错误", f"下载失败:\n{str(e)}", parent=download_window))
                    download_window.after(0, lambda: download_btn.config(state='normal'))
            
            threading.Thread(target=download_thread, daemon=True).start()
        
        download_btn = ttk.Button(
            button_frame,
            text="开始下载",
            command=start_download,
            bootstyle="success",
            width=12
        )
        download_btn.pack(side=tk.LEFT, padx=5)
        
        ttk.Button(
            button_frame,
            text="取消",
            command=download_window.destroy,
            bootstyle="secondary",
            width=12
        ).pack(side=tk.LEFT, padx=5)
    
    def _populate_model_options(self: 'MainWindow', frame, model_var, models_info, model_type):
        """填充模型选项"""
        # 清除现有选项
        for widget in frame.winfo_children():
            widget.destroy()
        
        if model_type == "rembg":
            # Rembg模型选项
            rembg_models = self.bg_remover.get_available_models()
            
            for model in rembg_models:
                rb = ttk.Radiobutton(
                    frame,
                    text=f"{model['name']} - {model['description']} ({model['size']}) - {model['best_for']}",
                    variable=model_var,
                    value=model['name']
                )
                rb.pack(anchor=tk.W, pady=2)
            
            # 添加安装说明
            note_label = ttk.Label(
                frame,
                text="\n💡 提示：Rembg需要先安装库：\n     pip install rembg\n     首次使用会自动下载模型。",
                font=("Arial", 9),
                foreground="gray"
            )
            note_label.pack(anchor=tk.W, pady=(10, 0))
            
        elif model_type == "sam":
            # SAM模型选项
            for key, info in models_info.items():
                rb = ttk.Radiobutton(
                    frame,
                    text=f"{key.upper()} - {info['description']} ({info['size']})",
                    variable=model_var,
                    value=key
                )
                rb.pack(anchor=tk.W, pady=2)
        else:
            # MiDaS模型选项
            midas_info = {
                'DPT_Large': {
                    'description': '最大、最准确（推荐）',
                    'size': '约1.4GB'
                },
                'DPT_Hybrid': {
                    'description': '中等大小',
                    'size': '约700MB'
                },
                'MiDaS_small': {
                    'description': '最小、最快',
                    'size': '约100MB'
                }
            }
            
            for key, info in midas_info.items():
                rb = ttk.Radiobutton(
                    frame,
                    text=f"{key} - {info['description']} ({info['size']})",
                    variable=model_var,
                    value=key
                )
                rb.pack(anchor=tk.W, pady=2)
            
            # 添加说明
            note_label = ttk.Label(
                frame,
                text="\n💡 提示：MiDaS模型可以手动下载，也可以在首次使用时自动下载。\n     手动下载会保存到torch缓存目录，与自动下载效果相同。",
                font=("Arial", 9),
                foreground="gray"
            )
            note_label.pack(anchor=tk.W, pady=(10, 0))
    
    def _update_model_selection(self: 'MainWindow', model_type, model_var, models_frame):
        """更新模型选择列表"""
        if model_type == "sam":
            sam_models_info = self.config.get('model.download_urls', {})
            model_var.set("vit_b")
            self._populate_model_options(models_frame, model_var, sam_models_info, "sam")
        elif model_type == "midas":
            model_var.set("DPT_Large")
            self._populate_model_options(models_frame, model_var, {}, "midas")
        else:  # rembg
            model_var.set("u2net")
            self._populate_model_options(models_frame, model_var, {}, "rembg")

